ARG IMAGE_NAME
FROM ${IMAGE_NAME}:{{ cuda.version.release_label }}-runtime-{{ cuda.os.distro }}{{ cuda.os.version }}{{ cuda.image_tag_suffix if "image_tag_suffix" in cuda }} as base

ENV NV_CUDA_LIB_VERSION "{{ cuda.version.release_label + "-1" }}"

{% if "x86_64" in cuda %}
FROM base as base-amd64

ENV NV_CUDA_CUDART_DEV_VERSION {{ cuda.x86_64.components.cudart_dev.version }}
ENV NV_NVML_DEV_VERSION {{ cuda.x86_64.components.nvml_dev.version }}
ENV NV_LIBCUSPARSE_DEV_VERSION {{ cuda.x86_64.components.libcusparse_dev.version }}
{% if "libnpp_dev" in cuda.x86_64.components %}
    {% set has_libnpp_dev_package = true %}
ENV NV_LIBNPP_DEV_VERSION {{ cuda.x86_64.components.libnpp_dev.version }}
ENV NV_LIBNPP_DEV_PACKAGE libnpp-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBNPP_DEV_VERSION}
{% endif %}

ENV NV_LIBCUBLAS_DEV_VERSION {{ cuda.x86_64.components.libcublas_dev.version }}
ENV NV_LIBCUBLAS_DEV_PACKAGE_NAME libcublas-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}
ENV NV_LIBCUBLAS_DEV_PACKAGE ${NV_LIBCUBLAS_DEV_PACKAGE_NAME}=${NV_LIBCUBLAS_DEV_VERSION}

{% if "libnccl2" in cuda.x86_64.components and cuda.x86_64.components.libnccl2 and "libnccl2_dev" in cuda.x86_64.components and cuda.x86_64.components.libnccl2_dev %}
    {% set has_libnccl_dev_package = true %}
ENV NV_LIBNCCL_DEV_PACKAGE_NAME libnccl-dev
ENV NV_LIBNCCL_DEV_PACKAGE_VERSION {{ cuda.x86_64.components.libnccl2_dev.version }}
ENV NCCL_VERSION {{ cuda.x86_64.components.libnccl2_dev.version }}
ENV NV_LIBNCCL_DEV_PACKAGE ${NV_LIBNCCL_DEV_PACKAGE_NAME}=${NV_LIBNCCL_DEV_PACKAGE_VERSION}+cuda{{ cuda.version.major }}.{{ cuda.version.minor }}

{% endif -%}

{% endif %}
{% if "arm64" in cuda %}
FROM base as base-arm64

ENV NV_CUDA_CUDART_DEV_VERSION {{ cuda.arm64.components.cudart_dev.version }}
ENV NV_NVML_DEV_VERSION {{ cuda.arm64.components.nvml_dev.version }}
ENV NV_LIBCUSPARSE_DEV_VERSION {{ cuda.arm64.components.libcusparse_dev.version }}
{% if "libnpp_dev" in cuda.arm64.components %}
    {% set has_libnpp_dev_package = true %}
ENV NV_LIBNPP_DEV_VERSION {{ cuda.arm64.components.libnpp_dev.version }}
ENV NV_LIBNPP_DEV_PACKAGE libnpp-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBNPP_DEV_VERSION}
{% endif %}

{% if cuda.flavor == "l4t" %}
    {% if cuda.version.major | int == 10 %}
ENV NV_LIBCUBLAS_DEV_PACKAGE_NAME libcublas-dev
    {% else %}
ENV NV_LIBCUBLAS_DEV_PACKAGE_NAME libcublas-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}
    {% endif %}
{% else %}
ENV NV_LIBCUBLAS_DEV_PACKAGE_NAME libcublas-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}
{% endif %}
ENV NV_LIBCUBLAS_DEV_VERSION {{ cuda.arm64.components.libcublas_dev.version }}
ENV NV_LIBCUBLAS_DEV_PACKAGE ${NV_LIBCUBLAS_DEV_PACKAGE_NAME}=${NV_LIBCUBLAS_DEV_VERSION}

{% if "libnccl2" in cuda.arm64.components and cuda.arm64.components.libnccl2 and "libnccl2_dev" in cuda.arm64.components and cuda.arm64.components.libnccl2_dev %}
    {% set has_libnccl_dev_package = true %}
ENV NV_LIBNCCL_DEV_PACKAGE_NAME libnccl-dev
ENV NV_LIBNCCL_DEV_PACKAGE_VERSION {{ cuda.arm64.components.libnccl2_dev.version }}
ENV NCCL_VERSION {{ cuda.arm64.components.libnccl2_dev.version }}
ENV NV_LIBNCCL_DEV_PACKAGE ${NV_LIBNCCL_DEV_PACKAGE_NAME}=${NV_LIBNCCL_DEV_PACKAGE_VERSION}+cuda{{ cuda.version.major }}.{{ cuda.version.minor }}

{% endif -%}
{% endif -%}

{% if "ppc64le" in cuda %}
FROM base as base-ppc64le

ENV NV_CUDA_CUDART_DEV_VERSION {{ cuda.ppc64le.components.cudart_dev.version }}
ENV NV_NVML_DEV_VERSION {{ cuda.ppc64le.components.nvml_dev.version }}
ENV NV_LIBCUSPARSE_DEV_VERSION {{ cuda.ppc64le.components.libcusparse_dev.version }}
{% if "libnpp_dev" in cuda.ppc64le.components %}
    {% set has_libnpp_dev_package = true %}
ENV NV_LIBNPP_DEV_VERSION {{ cuda.ppc64le.components.libnpp_dev.version }}
ENV NV_LIBNPP_DEV_PACKAGE libnpp-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBNPP_DEV_VERSION}
{% endif %}

ENV NV_LIBCUBLAS_DEV_VERSION {{ cuda.ppc64le.components.libcublas_dev.version }}
ENV NV_LIBCUBLAS_DEV_PACKAGE_NAME libcublas-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}
ENV NV_LIBCUBLAS_DEV_PACKAGE ${NV_LIBCUBLAS_DEV_PACKAGE_NAME}=${NV_LIBCUBLAS_DEV_VERSION}

{% if "libnccl2" in cuda.ppc64le.components and cuda.ppc64le.components.libnccl2 and "libnccl2_dev" in cuda.ppc64le.components and cuda.ppc64le.components.libnccl2_dev %}
    {% set has_libnccl_dev_package = true %}
ENV NV_LIBNCCL_DEV_PACKAGE_NAME libnccl-dev
ENV NV_LIBNCCL_DEV_PACKAGE_VERSION {{ cuda.ppc64le.components.libnccl2_dev.version }}
ENV NCCL_VERSION {{ cuda.ppc64le.components.libnccl2_dev.version }}
ENV NV_LIBNCCL_DEV_PACKAGE ${NV_LIBNCCL_DEV_PACKAGE_NAME}=${NV_LIBNCCL_DEV_PACKAGE_VERSION}+cuda{{ cuda.version.major }}.{{ cuda.version.minor }}

{% endif -%}

{% endif -%}

FROM base-${TARGETARCH}

ARG TARGETARCH

LABEL maintainer "NVIDIA CORPORATION <cudatools@nvidia.com>"

RUN apt-get update && apt-get install -y --no-install-recommends \
{% if cuda.os.version == "20.04" %}
    libtinfo5 libncursesw5 \
{% endif %}
    cuda-cudart-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_CUDA_CUDART_DEV_VERSION} \
    cuda-command-line-tools-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_CUDA_LIB_VERSION} \
    cuda-minimal-build-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_CUDA_LIB_VERSION} \
    cuda-libraries-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_CUDA_LIB_VERSION} \
    cuda-nvml-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_NVML_DEV_VERSION} \
    {% if has_libnpp_dev_package %}
    ${NV_LIBNPP_DEV_PACKAGE} \
    {% endif %}
    libcusparse-dev-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBCUSPARSE_DEV_VERSION} \
    ${NV_LIBCUBLAS_DEV_PACKAGE} \
    {% if has_libnccl_dev_package %}
    ${NV_LIBNCCL_DEV_PACKAGE} \
    {% endif %}
    && rm -rf /var/lib/apt/lists/*

# Keep apt from auto upgrading the cublas and nccl packages. See https://gitlab.com/nvidia/container-images/cuda/-/issues/88
RUN apt-mark hold ${NV_LIBCUBLAS_DEV_PACKAGE_NAME}{% if has_libnccl_dev_package %} ${NV_LIBNCCL_DEV_PACKAGE_NAME}{% endif %}


ENV LIBRARY_PATH /usr/local/cuda/lib64/stubs
